package com.keith.jpa.tenant;

import com.google.common.base.Strings;
import com.keith.jpa.tenant.context.TenantContextHolder;
import com.keith.jpa.tenant.utils.TenantUtils;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.*;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.*;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.Statements;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.replace.Replace;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import org.hibernate.resource.jdbc.spi.StatementInspector;
import org.springframework.util.CollectionUtils;

import java.util.Iterator;
import java.util.List;

/**
 * TenantJpaInspector.java
 * 拦截执行sql进行租户处理
 *
 * @author czykeih
 * @version 1.0 2023/10/25
 */
@Slf4j
public class TenantJpaInspector implements StatementInspector {
    @Override
    public String inspect(String s) {
        String retSql = s;
        if (TenantUtils.jpaConfig().isEnableTenant() && TenantUtils.jpaConfig().isLog()) {
            log.info("租户配置verifyTenant:[{}],tenantId:[{}]", TenantUtils.verifyTenant(), TenantContextHolder.getTenantId());
        }
        if (TenantUtils.checkCurrentTenantStatus()) {
            try {
                if (TenantUtils.jpaConfig().isLog()) {
                    log.info("解析前 SQL:[{}]", s);
                }
                long st = System.currentTimeMillis();
                Statements statements = CCJSqlParserUtil.parseStatements(s);
                StringBuilder newSql = new StringBuilder();
                int i = 0;
                for (Iterator<Statement> statementIterator = statements.getStatements().iterator(); statementIterator.hasNext(); ++i) {
                    Statement statement =statementIterator.next();
                    if (i > 0) {
                        newSql.append(";");
                    }
                    newSql.append(this.processParser(statement));
                }
                retSql = newSql.toString();
                if (TenantUtils.jpaConfig().isLog()) {
                    log.info("耗时:[{}],解析后 SQL:[{}]", System.currentTimeMillis() - st, retSql);
                }
            } catch (JSQLParserException e) {
                log.error("TenantJpaInspector Exception:", e);
                retSql = s;
            }
        }
        return retSql;
    }

    /**
     * 解析Sql
     *
     * @param statement
     * @return
     */
    private String processParser(Statement statement) {
        if (statement instanceof Insert) {
            this.processInsert((Insert) statement);
        } else if (statement instanceof Select) {
            this.processSelect((Select) statement);
        } else if (statement instanceof Update) {
            this.processUpdate((Update) statement);
        } else if (statement instanceof Delete) {
            this.processDelete((Delete) statement);
        } else if (statement instanceof Replace) {
            this.processReplace((Replace) statement);
        }
        return statement.toString();
    }

    /**
     * 解析删除
     *
     * @param delete
     */
    private void processDelete(Delete delete) {
        if (TenantUtils.verifyTableName(delete.getTable().getFullyQualifiedName())) {
            delete.setWhere(addExpression(delete.getTable(), delete.getWhere()));
        }
    }

    /**
     * 解析更新
     *
     * @param update
     */
    private void processUpdate(Update update) {
        if (TenantUtils.verifyTableName(update.getTable().getFullyQualifiedName())) {
            update.setWhere(addExpression(update.getTable(), update.getWhere()));
        }
    }

    /**
     * 解析查询
     *
     * @param select
     */
    private void processSelect(Select select) {
        processSelectBody(select.getSelectBody());
    }

    /**
     * 添加字段信息
     *
     * @param columns
     * @return
     */
    private int processInsertOrMergeColumn(List<Column> columns) {
        int index;
        if ((index = TenantUtils.verifyColumnName(columns)) < 0) {
            columns.add(new Column(TenantUtils.jpaConfig().getTenantColumn()));
        }
        return index;
    }

    /**
     * 判断租户字段是否有值
     *
     * @param index
     * @param itemsList
     */
    private void processInsertOrMergeParams(int index, ItemsList itemsList) {
        if (itemsList instanceof MultiExpressionList) {
            ((MultiExpressionList) itemsList).getExpressionLists().forEach(v -> {
                if (index < 0) {
                    v.getExpressions().add(TenantUtils.getTenantExpression());
                } else {
                    Expression expression = v.getExpressions().get(index);
                    if (expression instanceof NullValue) {
                        v.getExpressions().set(index, TenantUtils.getTenantExpression());
                    }
                }
            });
        } else {
            if (index < 0) {
                ((ExpressionList) itemsList).getExpressions().add(TenantUtils.getTenantExpression());
            } else {
                Expression expression = ((ExpressionList) itemsList).getExpressions().get(index);
                if (expression instanceof NullValue) {
                    ((ExpressionList) itemsList).getExpressions().set(index, TenantUtils.getTenantExpression());
                } else if (expression instanceof StringValue) {
                    String value = ((StringValue) expression).getValue();
                    if (Strings.isNullOrEmpty(value)) {
                        ((ExpressionList) itemsList).getExpressions().set(index, TenantUtils.getTenantExpression());
                    } else if (TenantUtils.verifyTenantIdValue(value)) {
                        ((ExpressionList) itemsList).getExpressions().set(index, TenantUtils.getTenantExpression());
                    }
                }
            }
        }
    }

    private void processReplace(Replace replace) {
        if (TenantUtils.verifyTableName(replace.getTable().getFullyQualifiedName())) {
            int index = processInsertOrMergeColumn(replace.getColumns());
            if (replace.getItemsList() != null) {
                processInsertOrMergeParams(index, replace.getItemsList());
            }
        }
    }

    /**
     * 解析插入
     */
    private void processInsert(Insert insert) {
        if (TenantUtils.verifyTableName(insert.getTable().getFullyQualifiedName())) {
            int index = processInsertOrMergeColumn(insert.getColumns());
            if (insert.getSelect() != null) {
                processPlainSelect((PlainSelect) insert.getSelect().getSelectBody(), true);
            } else if (insert.getItemsList() != null) {
                processInsertOrMergeParams(index, insert.getItemsList());
            }
        }
    }

    /**
     * 解析 plainSelect
     */
    private void processPlainSelect(PlainSelect plainSelect, boolean isAddColumn) {
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table table = (Table) fromItem;
            if (TenantUtils.verifyTableName(table.getFullyQualifiedName())) {
                plainSelect.setWhere(expression(plainSelect.getWhere(), table));
                if (isAddColumn) {
                    plainSelect.getSelectItems().add(new SelectExpressionItem(new Column(TenantUtils.jpaConfig().getTenantColumn())));
                }
            }
        } else {
            processFromItem(fromItem);
        }
        List<Join> joins = plainSelect.getJoins();
        if (!CollectionUtils.isEmpty(joins)) {
            for (Join join : joins) {
                processJoin(join);
                processFromItem(fromItem);
            }
        }
    }

    /**
     * 处理join查询
     *
     * @param join
     */
    private void processJoin(Join join) {
        if (join.getRightItem() instanceof Table) {
            Table table = (Table) join.getRightItem();
            if (TenantUtils.verifyTableName(table.getFullyQualifiedName())) {
                join.setOnExpression(expression(join.getOnExpression(), table));
            }
        }
    }

    /**
     * 针对update操作的条件处理
     *
     * @return
     */
    private Expression addExpression(Table table, Expression expression) {
        EqualsTo tenantExpression = new EqualsTo();
        tenantExpression.setLeftExpression(getAliasColumn(table));
        tenantExpression.setRightExpression(TenantUtils.getTenantExpression());
        if (expression != null) {
            if (expression instanceof OrExpression) {
                return new AndExpression(tenantExpression, new Parenthesis(expression));
            } else {
                return new AndExpression(tenantExpression, expression);
            }
        }
        return tenantExpression;
    }

    private Expression addWhereExpression(Table table){
        if(TenantUtils.verifyPenetration()){
            ExpressionList tenantIds = TenantUtils.getTenantIds(TenantContextHolder.getTenantId());
            return new InExpression(getAliasColumn(table), tenantIds);
        }
        return addExpression(table, null);
    }

    /**
     * 构建条件方法
     *
     * @param where
     * @param table
     * @return
     */
    private Expression expression(Expression where, Table table) {
        Expression binaryExpression = addWhereExpression(table);
        if (where == null) {
            return binaryExpression;
        }
        if (where instanceof BinaryExpression) {
            BinaryExpression binary = (BinaryExpression) where;
            expression(binary.getLeftExpression());
            expression(binary.getRightExpression());
        } else if (where instanceof InExpression) {
            expression(where);
        }
        if (where instanceof OrExpression) {
            return new AndExpression(new Parenthesis(where), binaryExpression);
        } else {
            return new AndExpression(where, binaryExpression);
        }
    }

    /**
     * 获取别名
     *
     * @param table
     * @return
     */
    private Column getAliasColumn(Table table) {
        StringBuilder column = new StringBuilder();
        if (table.getAlias() == null) {
            column.append(TenantUtils.jpaConfig().getTenantColumn());
        } else {
            column.append(table.getAlias().getName()).append(".").append(TenantUtils.jpaConfig().getTenantColumn());
        }
        return new Column(column.toString());
    }

    /**
     * 解析条件
     *
     * @param where
     */
    private void expression(Expression where) {
        if (where instanceof FromItem) {
            processFromItem((FromItem) where);
        } else if (where instanceof InExpression) {
            InExpression inExpression = (InExpression) where;
            ItemsList rightItemsList = inExpression.getRightItemsList();
            if (rightItemsList instanceof SubSelect) {
                processSelectBody(((SubSelect) rightItemsList).getSelectBody());
            }
        }
    }

    /**
     * 解析查询
     *
     * @param selectBody
     */
    private void processSelectBody(SelectBody selectBody) {
        if (selectBody != null) {
            if (selectBody instanceof PlainSelect) {
                this.processPlainSelect((PlainSelect) selectBody, false);
            } else if (selectBody instanceof WithItem) {
                WithItem withItem = (WithItem) selectBody;
                this.processSelectBody(withItem.getSelectBody());
            } else {
                SetOperationList operationList = (SetOperationList) selectBody;
                if (operationList.getSelects() != null && !operationList.getSelects().isEmpty()) {
                    operationList.getSelects().forEach(this::processSelectBody);
                }
            }
        }
    }

    /**
     * 处理sql子查询
     *
     * @param fromItem
     */
    private void processFromItem(FromItem fromItem) {
        if (fromItem instanceof SubJoin) {
            SubJoin subJoin = (SubJoin) fromItem;
            if (subJoin.getJoinList() != null) {
                subJoin.getJoinList().forEach(this::processJoin);
            }

            if (subJoin.getLeft() != null) {
                this.processFromItem(subJoin.getLeft());
            }
        } else if (fromItem instanceof SubSelect) {
            SubSelect subSelect = (SubSelect) fromItem;
            if (subSelect.getSelectBody() != null) {
                this.processSelectBody(subSelect.getSelectBody());
            }
        } else if (fromItem instanceof ValuesList) {
            log.debug("执行子查询[{}]", fromItem);
        } else if (fromItem instanceof LateralSubSelect) {
            LateralSubSelect lateralSubSelect = (LateralSubSelect) fromItem;
            if (lateralSubSelect.getSubSelect() != null) {
                SubSelect subSelect = lateralSubSelect.getSubSelect();
                if (subSelect.getSelectBody() != null) {
                    this.processSelectBody(subSelect.getSelectBody());
                }
            }
        }
    }

}
